package cc.git.liuyan.customeraiagent.core.provider.giteeai.embeddingmodel.bgem3;

import cc.git.liuyan.customeraiagent.core.embeddingmodel.EmbeddingModelInputData;
import cc.git.liuyan.customeraiagent.core.embeddingmodel.EmbeddingModelOutputData;
import cc.git.liuyan.customeraiagent.core.embeddingmodel.abstrct.AbstractEmbeddingModel;
import cc.git.liuyan.customeraiagent.core.provider.giteeai.embeddingmodel.GiteeAiEmbeddingModelApiClient;
import cc.git.liuyan.customeraiagent.core.provider.giteeai.embeddingmodel.api.GiteeAiEmbeddingModelApi;
import cc.git.liuyan.customeraiagent.core.provider.giteeai.embeddingmodel.api.reqandres.GiteeAiEmbeddingModelApiRequest;
import cc.git.liuyan.customeraiagent.core.provider.giteeai.embeddingmodel.api.reqandres.GiteeAiEmbeddingModelApiResponse;
import lombok.extern.slf4j.Slf4j;
import retrofit2.Call;
import retrofit2.Response;

import java.io.IOException;
import java.math.BigDecimal;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

//defauly sparse
@Slf4j
public class BgeM3GiteeAiEmbeddingModel extends AbstractEmbeddingModel {

    public BgeM3GiteeAiEmbeddingModel(String name, int dimension, GiteeAiEmbeddingModelApiClient apiClient) {
        super(name, dimension, apiClient);
    }

    @Override
    public List<BigDecimal> embed(String text) {
        GiteeAiEmbeddingModelApiRequest request = new GiteeAiEmbeddingModelApiRequest();
        request.setModel(this.modelName());
        request.setInput(text);

        GiteeAiEmbeddingModelApi api = (GiteeAiEmbeddingModelApi) this.apiClient.getApi();
        Call<GiteeAiEmbeddingModelApiResponse> responseCall = api.embeddings(new HashMap<>(), request);
        try {
            Response<GiteeAiEmbeddingModelApiResponse> execute = responseCall.execute();
            if (execute.isSuccessful()) {
                if (execute.body() != null) {
                    GiteeAiEmbeddingModelApiResponse response = execute.body();
                    List<BigDecimal> embedding = response.getData().get(0).getEmbedding();
                    return embedding;
                }
            } else {
                log.error(execute.message());
                throw new IOException(execute.message());
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        return null;
    }

    @Override
    public List<EmbeddingModelOutputData> embed(EmbeddingModelInputData inputData) {
        return super.embed(inputData);
    }

    @Override
    protected List<EmbeddingModelOutputData> embedFinal(List<EmbeddingModelInputData> inputDataSegmentList) {
        List<EmbeddingModelOutputData> allSegmentVectos = new ArrayList<>();
        for (EmbeddingModelInputData segmentInputData : inputDataSegmentList) {
            List<BigDecimal> embed = this.embed(segmentInputData.getInputData());
            EmbeddingModelOutputData embeddingModelOutputData = new EmbeddingModelOutputData(segmentInputData.getInputData());
            embeddingModelOutputData.setDenseVectors(embed);
            embeddingModelOutputData.setSparseVectors(embed);

            allSegmentVectos.add(embeddingModelOutputData);
        }
        return allSegmentVectos;
    }
}
